from langchain.prompts.example_selector.base import BaseExampleSelector
from typing import Dict, List
import numpy as np


class CustomExampleSelector(BaseExampleSelector):

    def __init__(self, examples: List[Dict[str, str]]):
        self.examples = examples

    def add_example(self, example: Dict[str, str]) -> None:
        """存储的新示例"""
        self.examples.append(example)

    def select_examples(self, input_variables: Dict[str, str]) -> List[dict]:
        """根据输入选择要使用的示例"""
        return np.random.choice(self.examples, size=2, replace=False)


examples = [
    {"foo": "1"},
    {"foo": "2"},
    {"foo": "3"}
]

# 初始化选择器
example_selector = CustomExampleSelector(examples)

# 随机选择2个
example_selector.select_examples({"foo": "foo"})
print(example_selector)

# 添加一个新的示例
example_selector.add_example({"foo": "4"})

# 随机选择2个
example_selector.select_examples({"foo": "foo"})
